from collections import OrderedDict
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import pickle
import math

from sklearn.base import ClassifierMixin

from pre_process.read_data import read_to_predict_data, DATA_TRAIN_COLUMN

def get_daily_pnl(data_dir: str, model_dir: str, symbol: str, model_name: str):
    with open(model_dir, mode='rb') as f:
        model = pickle.load(f)
    assert(isinstance(model, ClassifierMixin))
    factor, _ = read_to_predict_data(data_dir)
    predict_direction = model.predict(factor[DATA_TRAIN_COLUMN].reset_index())
    # print(predict_direction)
    factor['next change'] = pd.DataFrame(predict_direction)
    n_bar = len(factor)
    signal = pd.Series([0] * n_bar)
    signal[predict_direction > 0] = 1
    signal[predict_direction < 0] = -1
    position_pos = pd.Series([np.nan] * n_bar)
    position_pos[0] = 0
    position_pos[signal == 1] = 1
    position_pos[signal == -1] = 0
    position_pos.ffill(inplace=True)
    position_neg = pd.Series([np.nan] * n_bar)
    position_neg[0] = 0
    position_neg[signal == -1] = -1
    position_neg[signal == 1] = 0
    position_neg.ffill(inplace=True)
    position = position_pos + position_neg
    position[0]=0
    position[n_bar - 1] = 0
    position[n_bar - 2] = 0
    change_pos = position - position.shift(1)
    change_pos[0] = 0
    change_base = pd.Series([0] * n_bar)
    change_buy = change_pos > 0
    change_sell = change_pos < 0
    change_base[change_buy] = factor["close"][change_buy]
    change_base[change_sell] = factor["close"][change_sell]
    # pos = change_pos.copy()
    # pos[pos > 1] = 1
    # pos[pos < -1] = -1
    # pnl = change_base * pos 
    # pnl[pnl == 0] = np.nan
    # pnl[0] = 0
    # pnl.bfill(inplace=True)
    # pnl.fillna(0)
    # pnl[pnl > 0]
    pre_pos = 0 # empty holding
    profit = 0
    hold = 0
    profit_list = []
    for i in range(len(change_pos)):
        tmp_profit = profit

        if hold > 0:
            tmp_profit = tmp_profit + factor['close'].loc[i] - hold
        elif hold < 0:
            tmp_profit = tmp_profit + abs(hold) - factor['close'].loc[i]

        pos = change_pos.loc[i]
        if pos > 0:
            if pos == 1:
                if pre_pos == -1:
                    # 清仓
                    pre_pos = 0
                    hold = 0
                    profit = tmp_profit
                else:
                    hold = factor['close'].loc[i]
                    pre_pos = 0
            elif pos == 2:
                pre_pos = -1
                hold = factor['close'].loc[i]
                profit = tmp_profit
        elif pos < 0:
            if pos == -1:
                if pre_pos == 1:
                    # 清仓
                    pre_pos = 0
                    hold = 0
                    profit = tmp_profit
                else:
                    hold = -factor['close'].loc[i]
                    pre_pos = 0
            elif pos == -2:
                pre_pos = 1
                hold = -factor['close'].loc[i]
                profit = tmp_profit
        # print(tmp_profit)
        profit_list.append(tmp_profit)

    factor['pnl'] = pd.Series(profit_list)
    final_pnl = -sum(change_base * change_pos)
    # turnover = sum(change_base * abs(change_pos))
    # num = sum((position != 0) & (change_pos != 0))
    # hld_period = sum(position != 0)

    # result = OrderedDict([("final.pnl", final_pnl), ("turnover", turnover), ("num", num), ("hld.period", hld_period)])
    # factor.set_index('trade_date', inplace=True)
    # dates_date = factor.index
    date_format = [pd.to_datetime(str(d)) for d in factor['trade_date']]
    plt.clf()
    plt.cla()
    plt.figure(1, figsize=(16, 10))
    plt.suptitle("%s %s pnl (daily floating)" % (model_name, symbol))
    plt.title('fix profit %d' % final_pnl)
    plt.xlabel("date")
    plt.ylabel("pnl")
    # print(factor['pnl'])
    plt.plot(date_format, factor['pnl'])
    plt.savefig('./analysis/%s_%s.jpg' % (model_name, symbol))
